import torch.utils.data as data
from torchvision import transforms
from PIL import Image
import os
import torch
import numpy as np

from .util.mask import (bbox2mask, brush_stroke_mask, get_irregular_mask, random_bbox, random_cropping_bbox)

IMG_EXTENSIONS = [
    '.jpg', '.JPG', '.jpeg', '.JPEG',
    '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',
]

def is_image_file(filename):
    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)

def make_dataset(dir):
    if os.path.isfile(dir):
        images = [i for i in np.genfromtxt(dir, dtype=np.str, encoding='utf-8')]
    else:
        images = []
        assert os.path.isdir(dir), '%s is not a valid directory' % dir
        for root, _, fnames in sorted(os.walk(dir)):
            for fname in sorted(fnames):
                if is_image_file(fname):
                    path = os.path.join(root, fname)
                    images.append(path)

    return images

def pil_loader(path):
    return Image.open(path).convert('RGB')

class InpaintDataset(data.Dataset):
    def __init__(self, data_root, mask_config={}, data_len=-1, image_size=[256, 256], loader=pil_loader):
        imgs = make_dataset(data_root)
        if data_len > 0:
            self.imgs = imgs[:int(data_len)]
        else:
            self.imgs = imgs
        self.tfs = transforms.Compose([
                transforms.Resize((image_size[0], image_size[1])),
                transforms.ToTensor(),
                transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5,0.5, 0.5])
        ])
        self.loader = loader
        self.mask_config = mask_config
        self.mask_mode = self.mask_config['mask_mode']
        self.image_size = image_size

    def __getitem__(self, index):
        ret = {}
        path = self.imgs[index]
        img = self.tfs(self.loader(path))
        mask = self.get_mask()
        cond_image = img*(1. - mask) + mask*torch.randn_like(img)
        mask_img = img*(1. - mask) + mask

        ret['gt_image'] = img
        ret['cond_image'] = cond_image
        ret['mask_image'] = mask_img
        ret['mask'] = mask
        ret['path'] = path.rsplit("/")[-1].rsplit("\\")[-1]
        return ret

    def __len__(self):
        return len(self.imgs)

    def get_mask(self):
        if self.mask_mode == 'bbox':
            mask = bbox2mask(self.image_size, random_bbox())
        elif self.mask_mode == 'center':
            h, w = self.image_size
            mask = bbox2mask(self.image_size, (h//4, w//4, h//2, w//2))
        elif self.mask_mode == 'irregular':
            mask = get_irregular_mask(self.image_size)
        elif self.mask_mode == 'free_form':
            mask = brush_stroke_mask(self.image_size)
        elif self.mask_mode == 'hybrid':
            regular_mask = bbox2mask(self.image_size, random_bbox())
            irregular_mask = brush_stroke_mask(self.image_size, )
            mask = regular_mask | irregular_mask
        elif self.mask_mode == 'file':
            pass
        else:
            raise NotImplementedError(
                f'Mask mode {self.mask_mode} has not been implemented.')
        return torch.from_numpy(mask).permute(2,0,1)


class UncroppingDataset(data.Dataset):
    def __init__(self, data_root, mask_config={}, data_len=-1, image_size=[256, 256], loader=pil_loader):
        imgs = make_dataset(data_root)
        if data_len > 0:
            self.imgs = imgs[:int(data_len)]
        else:
            self.imgs = imgs
        self.tfs = transforms.Compose([
                transforms.Resize((image_size[0], image_size[1])),
                transforms.ToTensor(),
                transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5,0.5, 0.5])
        ])
        self.loader = loader
        self.mask_config = mask_config
        self.mask_mode = self.mask_config['mask_mode']
        self.image_size = image_size

    def __getitem__(self, index):
        ret = {}
        path = self.imgs[index]
        img = self.tfs(self.loader(path))
        mask = self.get_mask()
        cond_image = img*(1. - mask) + mask*torch.randn_like(img)
        mask_img = img*(1. - mask) + mask

        ret['gt_image'] = img
        ret['cond_image'] = cond_image
        ret['mask_image'] = mask_img
        ret['mask'] = mask
        ret['path'] = path.rsplit("/")[-1].rsplit("\\")[-1]
        return ret

    def __len__(self):
        return len(self.imgs)

    def get_mask(self):
        if self.mask_mode == 'manual':
            mask = bbox2mask(self.image_size, self.mask_config['shape'])
        elif self.mask_mode == 'fourdirection' or self.mask_mode == 'onedirection':
            mask = bbox2mask(self.image_size, random_cropping_bbox(mask_mode=self.mask_mode))
        elif self.mask_mode == 'hybrid':
            if np.random.randint(0,2)<1:
                mask = bbox2mask(self.image_size, random_cropping_bbox(mask_mode='onedirection'))
            else:
                mask = bbox2mask(self.image_size, random_cropping_bbox(mask_mode='fourdirection'))
        elif self.mask_mode == 'file':
            pass
        else:
            raise NotImplementedError(
                f'Mask mode {self.mask_mode} has not been implemented.')
        return torch.from_numpy(mask).permute(2,0,1)


class ColorizationDataset(data.Dataset):
    def __init__(self, data_root, data_flist, data_len=-1, image_size=[224, 224], loader=pil_loader):
        self.data_root = data_root
        # flist = make_dataset(data_flist)
        # if data_len > 0:
        #     self.flist = flist[:int(data_len)]
        # else:
        #     self.flist = flist
        self.tfs = transforms.Compose([
                transforms.Resize((image_size[0], image_size[1])),
                transforms.ToTensor(),
                transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5,0.5, 0.5])
        ])
        self.loader = loader
        self.image_size = image_size


        clean_files = sorted(os.listdir(os.path.join(data_root, 'gt')))
        # print(clean_files)
        noisy_files = sorted(os.listdir(os.path.join(data_root, 'input')))
        # print(noisy_files)
        self.clean_filenames = [os.path.join(data_root, 'gt', x) for x in clean_files ]
        self.noisy_filenames = [os.path.join(data_root, 'input', x)       for x in noisy_files ]
        self.tar_size = len(self.clean_filenames)  # get the size of target
        
    def __getitem__(self, index):
        ret = {}
        # file_name = str(self.flist[index]).zfill(5) + '.jpg'
        tar_index   = index % self.tar_size
        
        img = self.tfs(self.loader(self.clean_filenames[tar_index]))
        cond_image = self.tfs(self.loader(self.noisy_filenames[tar_index]))


        clean_filename = os.path.split(self.clean_filenames[tar_index])[-1]
        noisy_filename = os.path.split(self.noisy_filenames[tar_index])[-1]

        ret['gt_image'] = img
        ret['cond_image'] = cond_image
        ret['path'] = clean_filename
        return ret

    def __len__(self):
        return self.tar_size


# class MyDataset(Dataset):
#     def __init__(self, rgb_dir):
#         # self.data = [
#         #     torch.randn(256, 256, 3),
#         #     torch.randn(256, 256, 3),
#         # ]
#         # self.target_transform = target_transform

#         clean_files = sorted(os.listdir(os.path.join(rgb_dir, 'gt')))
#         # print(clean_files)
#         noisy_files = sorted(os.listdir(os.path.join(rgb_dir, 'input')))
#         # print(noisy_files)
#         self.clean_filenames = [os.path.join(rgb_dir, 'gt', x) for x in clean_files ]
#         self.noisy_filenames = [os.path.join(rgb_dir, 'input', x)       for x in noisy_files ]
#         # print(self.clean_filenames, self.noisy_filenames)
#         # self.img_options=img_options

#         self.tar_size = len(self.clean_filenames)  # get the size of target

#     def __getitem__(self, index):
#         tar_index   = index % self.tar_size
#         clean = torch.from_numpy(np.float32(load_img(self.clean_filenames[tar_index])))
#         noisy = torch.from_numpy(np.float32(load_img(self.noisy_filenames[tar_index])))
        
#         clean = clean.permute(2,0,1)
#         noisy = noisy.permute(2,0,1)

#         clean_filename = os.path.split(self.clean_filenames[tar_index])[-1]
#         noisy_filename = os.path.split(self.noisy_filenames[tar_index])[-1]

#         #Crop Input and Target
#         ps = 256
#         H = clean.shape[1]
#         W = clean.shape[2]
#         r = np.random.randint(0, H - ps)
#         c = np.random.randint(0, W - ps)
#         clean = clean[:, r:r + ps, c:c + ps]
#         noisy = noisy[:, r:r + ps, c:c + ps]

#         noisy = noisy.permute(1, 2, 0)
#         clean = clean.permute(1, 2, 0)
#         # apply_trans = transforms_aug[random.getrandbits(3)]

#         # clean = getattr(augment, apply_trans)(clean)
#         # noisy = getattr(augment, apply_trans)(noisy)   
#         # return {
#         #     "image": clean,   #gt
#         #     "lq_image": noisy, #input
#         # }
#         ret['gt_image'] = clean
#         ret['cond_image'] = noisy
#         ret['path'] = file_name
#         return ret

#     def __len__(self):
#         return self.tar_size
        
#         # dataset = MyDataset(opt.train_path)